import os
import base64
import hashlib
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.primitives import padding
from cryptography.hazmat.backends import default_backend
import rsa

class CustomCryptoSystem:
    def __init__(self):
        self.backend = default_backend()
        
    def derive_key_from_password(self, password, salt=None, key_length=32):
        """从用户密码派生加密密钥"""
        if salt is None:
            salt = os.urandom(16)  # 生成16字节的随机盐值
        
        # 使用PBKDF2算法从密码派生密钥
        # 这里为了简化使用hashlib，实际应用中应使用更安全的密钥派生函数
        dk = hashlib.pbkdf2_hmac(
            'sha256',  # 使用SHA-256哈希算法
            password.encode('utf-8'),  # 密码转字节
            salt,  # 盐值
            100000,  # 迭代次数
            dklen=key_length  # 密钥长度
        )
        
        return dk, salt
    
    def encrypt_file_with_password(self, file_path, password, output_path=None):
        """使用用户自定义密码加密文件"""
        try:
            # 生成随机IV
            iv = os.urandom(16)
            
            # 从密码派生密钥
            key, salt = self.derive_key_from_password(password)
            
            # 读取文件内容
            with open(file_path, 'rb') as file:
                file_data = file.read()
            
            # 创建加密器
            cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=self.backend)
            encryptor = cipher.encryptor()
            
            # 对数据进行填充
            padder = padding.PKCS7(algorithms.AES.block_size).padder()
            padded_data = padder.update(file_data) + padder.finalize()
            
            # 加密数据
            encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
            
            # 组合加密数据、盐值和IV
            result = salt + iv + encrypted_data
            
            # 生成输出路径
            if output_path is None:
                output_path = file_path + '.encrypted'
            
            # 写入加密文件
            with open(output_path, 'wb') as file:
                file.write(result)
            
            return output_path
        except Exception as e:
            raise Exception(f"文件加密失败: {str(e)}")
    
    def decrypt_file_with_password(self, file_path, password, output_path=None):
        """使用用户自定义密码解密文件"""
        try:
            # 读取加密文件内容
            with open(file_path, 'rb') as file:
                encrypted_data = file.read()
            
            # 提取盐值、IV和加密数据
            salt = encrypted_data[:16]
            iv = encrypted_data[16:32]
            cipher_data = encrypted_data[32:]
            
            # 从密码派生密钥
            key, _ = self.derive_key_from_password(password, salt)
            
            # 创建解密器
            cipher = Cipher(algorithms.AES(key), modes.CBC(iv), backend=self.backend)
            decryptor = cipher.decryptor()
            
            # 解密数据
            padded_data = decryptor.update(cipher_data) + decryptor.finalize()
            
            # 去除填充
            unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
            data = unpadder.update(padded_data) + unpadder.finalize()
            
            # 生成输出路径
            if output_path is None:
                output_path = file_path.replace('.encrypted', '')
                if os.path.exists(output_path):
                    output_path = output_path.split(".")[0] + "_decrypted." + output_path.split(".")[-1]
            
            # 写入解密文件
            with open(output_path, 'wb') as file:
                file.write(data)
            
            return output_path
        except Exception as e:
            raise Exception(f"文件解密失败: {str(e)}")
    
    def generate_rsa_keys_with_password(self, key_size=2048, password=None):
        """生成RSA密钥对，并可选择使用密码保护私钥"""
        try:
            # 生成RSA密钥对
            public_key, private_key = rsa.newkeys(key_size)
            
            return public_key, private_key
        except Exception as e:
            raise Exception(f"RSA密钥对生成失败: {str(e)}")
    
    def save_rsa_keys(self, public_key, private_key, public_path, private_path, password=None):
        """保存RSA密钥对到文件"""
        try:
            # 保存公钥
            with open(public_path, 'wb') as file:
                file.write(public_key.save_pkcs1())
            
            # 保存私钥，可选择使用密码保护
            with open(private_path, 'wb') as file:
                if password:
                    file.write(private_key.save_pkcs1(passphrase=password))
                else:
                    file.write(private_key.save_pkcs1())
            
            return True
        except Exception as e:
            raise Exception(f"RSA密钥保存失败: {str(e)}")
    
    def encrypt_file_with_rsa(self, file_path, public_key_path, output_path=None):
        """使用RSA公钥加密文件"""
        try:
            # 读取文件内容
            with open(file_path, 'rb') as file:
                file_content = file.read()
            
            # 读取公钥
            with open(public_key_path, 'rb') as key_file:
                public_key = rsa.PublicKey.load_pkcs1(key_file.read())
            
            # 注意：RSA只适合加密小数据，这里我们使用RSA加密AES密钥，然后用AES加密文件
            # 生成随机AES密钥
            aes_key = os.urandom(32)  # 256位AES密钥
            
            # 用RSA加密AES密钥
            encrypted_aes_key = rsa.encrypt(aes_key, public_key)
            
            # 生成随机IV
            iv = os.urandom(16)
            
            # 创建AES加密器
            cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=self.backend)
            encryptor = cipher.encryptor()
            
            # 对文件数据进行填充
            padder = padding.PKCS7(algorithms.AES.block_size).padder()
            padded_data = padder.update(file_content) + padder.finalize()
            
            # 用AES加密文件数据
            encrypted_data = encryptor.update(padded_data) + encryptor.finalize()
            
            # 组合加密后的AES密钥、IV和加密的文件数据
            result = encrypted_aes_key + iv + encrypted_data
            
            # 生成输出路径
            if output_path is None:
                output_path = file_path + '.rsa_encrypted'
            
            # 写入加密文件
            with open(output_path, 'wb') as file:
                file.write(result)
            
            return output_path
        except Exception as e:
            raise Exception(f"RSA文件加密失败: {str(e)}")
    
    def decrypt_file_with_rsa(self, file_path, private_key_path, password=None, output_path=None):
        """使用RSA私钥解密文件"""
        try:
            # 后缀
            suffix = file_path.split(".")[-1]
            # 读取加密文件内容
            with open(file_path, 'rb') as file:
                encrypted_data = file.read()
            
            # 读取私钥
            with open(private_key_path, 'rb') as key_file:
                if password:
                    private_key = rsa.PrivateKey.load_pkcs1(key_file.read(), passphrase=password)
                else:
                    private_key = rsa.PrivateKey.load_pkcs1(key_file.read())
            
            # 获取RSA密钥长度以确定加密的AES密钥的大小
            key_size_bytes = (private_key.n.bit_length() + 7) // 8
            
            # 提取加密的AES密钥、IV和加密的文件数据
            encrypted_aes_key = encrypted_data[:key_size_bytes]
            iv = encrypted_data[key_size_bytes:key_size_bytes + 16]
            cipher_data = encrypted_data[key_size_bytes + 16:]
            
            # 用RSA解密AES密钥
            aes_key = rsa.decrypt(encrypted_aes_key, private_key)
            
            # 创建AES解密器
            cipher = Cipher(algorithms.AES(aes_key), modes.CBC(iv), backend=self.backend)
            decryptor = cipher.decryptor()
            
            # 解密文件数据
            padded_data = decryptor.update(cipher_data) + decryptor.finalize()
            
            # 去除填充
            unpadder = padding.PKCS7(algorithms.AES.block_size).unpadder()
            data = unpadder.update(padded_data) + unpadder.finalize()
            
            # 生成输出路径
            if output_path is None:
                output_path = file_path.replace('.rsa_encrypted', '')
                if os.path.exists(output_path):
                    pl = output_path.split(".")
                    output_path = pl[0] + "_rsa_decrypted." + pl[-1]
            
            # 写入解密文件
            with open(output_path, 'wb') as file:
                file.write(data)
            
            return output_path
        except Exception as e:
            raise Exception(f"RSA文件解密失败: {str(e)}")

# 创建一个全局实例，方便使用
crypto_system = CustomCryptoSystem()